Greptile Summary

The PR introduces an exported wrapper for a cuDNN fused MLA Q up-projection, RoPE, and MXFP8 quantization operation, while extending fused attention to accept prequantized MXFP8 Q/K/V and optionally emit BF16 gradients.

• Adds the FusedMLAQUpProjRopeQuant public API and SM100 kernel capability detection.
• Adds split MXFP8 quantization and scale transpose/swizzle helpers.
• Extends fused-attention layout detection and backward output handling for prequantized MXFP8 tensors.

Confidence Score: 5/5

The PR appears safe to merge based on the reviewed changes, with no concrete blocking or independently actionable non-blocking defect established.

The new MXFP8 layout transformations follow the established fast-path conventions, the BF16-gradient option is consistently forwarded through direct and checkpointed fused-attention paths, and no reachable changed-code failure remains.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  X["Q low-rank activation"] --> F["Fused Q up-projection + RoPE"]
  W["BF16 or MXFP8 weight"] --> F
  C["RoPE cosine/sine"] --> F
  F --> Q["Prequantized MXFP8 Q"]
  K["K activation"] --> KV["MXFP8 quantization"]
  V["V activation"] --> KV
  Q --> S["BHSD scale transpose + swizzle"]
  KV --> S
  S --> A["cuDNN fused attention"]
  A --> B["BF16 dQ/dK/dV when requested"]

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}